/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Definition of the particle_set classes. \ingroup sfrhist

// $Id$

#ifndef __particle_set__
#define __particle_set__

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>
#include <cassert>
#include <exception>
#include "mcrx-types.h"


namespace mcrx {
  template<typename> class particle_ptr;
  class particle_set_generic;
  class gas_particle_set;
  class star_particle_set;
  class dm_particle_set;
  class bh_particle_set;
  class Snapshot;
  class sfrhist_snapshot;
  class stellarmodel;
};

namespace CCfits {
  class ExtHDU;
  class Column;
};

/** Simple container for "pointing" to a particle in a set. The type
    of particle set is an optional template parameter that can be used
    to make it refer to one of the specialized particle sets.  */
template <typename T_set=mcrx::particle_set_generic>
class mcrx::particle_ptr {
  T_set* set_;
  int num_;
public:
  particle_ptr(T_set& s, int n) : set_(&s), num_(n) {};
  T_set& set() const {return *set_;};
  int num() const {return num_;};
};


/** Base class for classes containing quantities for a set of
    particles. */
class mcrx::particle_set_generic {
  friend class mcrx::Snapshot;
  friend class mcrx::sfrhist_snapshot;
protected:
  /// Maps particle ID to  position in vector.
  std::map<unsigned int, int> order_; 
  /// ID number.
  std::vector<unsigned int> id_;
  /// Mass.
  std::vector<T_float> m_;
  /// Position.
  std::vector<vec3d> pos_;
  /// Velocity.
  std::vector<vec3d> vel_;
  /// Radius.
  std::vector<T_float> h_;

  /** Output the quantities for a certain particle in the set to the
      stream. */
  virtual void write_ascii_table_line(std::ostream&, int) const;

public:
  virtual ~particle_set_generic() {};

  /// Exception class thrown when a non-existent ID number is requested.
  class illegal_id : public std::exception {
  public:
    const char* what() const throw() {return "Illegal ID number";}
  };

  /// Returns index into particle vectors for a given ID. If the ID
  /// doesn't exist, it throws illegal_id.
  int find_ID(unsigned int id) const;

  int min_ID() const {assert(n()>0); return *std::min_element(id_.begin(), id_.end());};
  int max_ID() const {assert(n()>0); return *std::max_element(id_.begin(), id_.end());};

  /// Returns the number of particles in the set
  int n() const { return id_.size();};

  /// \name Access functions
  ///@{
  unsigned int& id(int i) {return id_[i];};
  const unsigned int& id(int i) const {return id_[i];};
  T_float& m(int i) {return m_[i];};
  const T_float& m(int i) const {return m_[i];};
  vec3d& pos(int i) {return pos_[i];};
  const vec3d& pos(int i) const {return pos_[i];};
  vec3d& vel(int i) {return vel_[i];};
  const vec3d& vel(int i) const {return vel_[i];};
  T_float& h(int i) {return h_[i];};
  const T_float& h(int i) const {return h_[i];};
  ///@}

  /** Convert units of quantities. The arguments supply the current
      units in terms of the units we want, ie if mass is 1.5 then the
      mass field will be multiplied by 1.5. */
  virtual void convert_units(T_float mass, T_float length, T_float time);

  /** Write particle data to a FITS table. The columns depend on the
      particle species and are predefined. If the columns for the
      species don't exist in the table, CCfits will throw an
      exception. */
  virtual void write_table_chunk (int, int, int, CCfits::ExtHDU&) const;

  /** Write particle data to an std::ostream. */
  void write_ascii_table (std::ostream&) const;

  /** Load particle set from PARTICLEDATA HDU. */
  virtual void load_particledata(CCfits::ExtHDU&);

  /** Concatenate two particle sets. There must be no duplicate
      IDs. */
  particle_set_generic operator+(const particle_set_generic& rhs) const {
    particle_set_generic p(*this); p+=rhs; return p; };
  /** Add a particle set to this. There must be no duplicate IDs. */
  particle_set_generic& operator+=(const particle_set_generic& rhs);
};


/** Data for sets of gas particles.  */
class mcrx::gas_particle_set : public mcrx::particle_set_generic {
protected:
public:
  std::vector<T_float> z_;  ///< Metallicity. (As a mass fraction.)
  std::vector<T_float> sfr_; ///< Star-formation rate.
  std::vector<T_float> temp_; ///< Gas temperature.
  std::vector<T_float> teff_; ///< Effective gas temperature including feedback.
  std::vector<unsigned int> pid_; ///< ID of parent gas particle if applicable.

  /** Output the quantities for a certain particle in the set to the
      stream. */
  virtual void write_ascii_table_line(std::ostream&, int) const;

public:
  virtual void convert_units(T_float mass, T_float length, T_float time);

  /// \name Access functions
  ///@{
  T_float& z(int i) {return z_[i];};
  const T_float& z(int i) const {return z_[i];};
  T_float& sfr(int i) {return sfr_[i];};
  const T_float& sfr(int i) const {return sfr_[i];};
  T_float& temp(int i) {return temp_[i];};
  const T_float& temp(int i) const {return temp_[i];};
  T_float& teff(int i) {return teff_[i];};
  const T_float& teff(int i) const {return teff_[i];};

  unsigned int& pid(int i) {return pid_[i];};
  const unsigned int& pid(int i) const {return pid_[i];};
  //@}

  virtual void write_table_chunk (int, int, int, CCfits::ExtHDU&) const;
  virtual void load_particledata(CCfits::ExtHDU&);

  /** Concatenate two particle sets. */
  gas_particle_set operator+(const gas_particle_set& rhs) const {
    gas_particle_set p(*this); p+=rhs; return p; };
  gas_particle_set& operator+=(const gas_particle_set& rhs);
};


/** Data for sets of star particles. This includes both stars created
    during the simulation and star particles present at the start. */
class mcrx::star_particle_set : public mcrx::particle_set_generic {
protected:
public:
  std::vector<T_float> z_;  ///< Metallicity. (As a mass fraction.)
  /// Formation time.
  std::vector<T_float> tf_; 
  /// ID of parent gas particle, if applicable.
  std::vector<unsigned int> pid_; 
  /** Mass at creation time. Since stellar populations recycle mass to
      the gas phase, this will in general be different from the
      current (gravitational) mass of the particle. */
  std::vector<T_float> mcr_; 

  /** Calculate the creation mass based on the current mass at the
      given time, based on the mass loss info in the stellar model. */
  void calculate_creation_mass(T_float time, const mcrx::stellarmodel& m);

  /** Output the quantities for a certain particle in the set to the
      stream. */
  virtual void write_ascii_table_line(std::ostream&, int) const;

public:
  virtual void convert_units(T_float mass, T_float length, T_float time);

  /// \name Access functions
  ///@{
  T_float& z(int i) {return z_[i];};
  const T_float& z(int i) const {return z_[i];};
  T_float& tf(int i) {return tf_[i];};
  const T_float& tf(int i) const {return tf_[i];};
  unsigned int& pid(int i) {return pid_[i];};
  const unsigned int& pid(int i) const {return pid_[i];};
  T_float& mcr(int i) {return mcr_[i];};
  const T_float& mcr(int i) const {return mcr_[i];};
  ///@}

  virtual void write_table_chunk (int, int, int, CCfits::ExtHDU&) const;
  virtual void load_particledata(CCfits::ExtHDU&);

  /** Concatenate two particle sets. */
  star_particle_set operator+(const star_particle_set& rhs) const {
    star_particle_set p(*this); p+=rhs; return p; };
  star_particle_set& operator+=(const star_particle_set& rhs);
};


/** Data for sets of dark-matter particles. Contains nothing but the
    generic quantities. */
class mcrx::dm_particle_set : public mcrx::particle_set_generic {

public:
  virtual void convert_units(T_float mass, T_float length, T_float time);

  virtual void write_table_chunk (int, int, int, CCfits::ExtHDU&) const;
  virtual void load_particledata(CCfits::ExtHDU&);

  /** Concatenate two particle sets. */
  dm_particle_set operator+(const dm_particle_set& rhs) const {
    dm_particle_set p(*this); p+=rhs; return p; };
  dm_particle_set& operator+=(const dm_particle_set& rhs);
};

/** Data for sets of black-hole particles. This includes BH mass
    (i.e., not including gas associated with BH particle) and BH mass
    accretion rate. */
class mcrx::bh_particle_set : public mcrx::particle_set_generic {

protected:
public:
  std::vector<T_float> mbh_; ///< Black hole mass.
  std::vector<T_float> mdotbh_; ///< Black hole accretion rate.

public:
  virtual void convert_units(T_float mass, T_float length, T_float time);

  /// \name Access functions
  ///@{
  T_float& mbh(int i) {return mbh_[i];};
  const T_float& mbh(int i) const {return mbh_[i];};
  T_float& mdotbh(int i) {return mdotbh_[i];};
  const T_float& mdotbh(int i) const {return mdotbh_[i];};
  ///@}

  virtual void write_table_chunk (int, int, int, CCfits::ExtHDU&) const;
  virtual void load_particledata(CCfits::ExtHDU&);

  /** Concatenate two particle sets. */
  bh_particle_set operator+(const bh_particle_set& rhs) const {
    bh_particle_set p(*this); p+=rhs; return p; };
  bh_particle_set& operator+=(const bh_particle_set& rhs);
};

#endif

